package com.jboltai.capability.model.baidu;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.capability.model.siliconflow.SiliconflowAIImpl;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ThinkStatus;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.util.http.HttpUtil;
import com.jboltai.util.http.RequestInstance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

public class WenxinAIImpl implements AIEventProcessor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public static final WenxinAIImpl INSTANCE = new WenxinAIImpl();

    private static final String STOP_FLAG = "[DONE]";

    public static final String API_URL = "https://qianfan.baidubce.com/v2/chat/completions";

    @Override
    public void process(AIChatEvent event, AIResource resource) {
        //构建请求实例
        RequestInstance request = buildRequest(event, resource);

        event.setAttr(EventAttrKey.REQUEST_INSTANCE,request);

        LOGGER.debug("开始 调用百度千帆大模型 处理AIChat事件：\n\n\nEvent参数: {}, \n\n\nRequest参数: {}", JSON.toJSONString(event), JSON.toJSONString(request));
        boolean stream = event.getIsStream();
        if (stream) {
            //流式请求
            processEventByStream(event, resource, request);

        } else {
            processEventByNormal(event, resource, request);
        }
    }

    /**
     * 通过正常请求 处理event
     * @param event
     * @param request
     */
    private static void processEventByNormal(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, false);
        request.asyncSend(response -> {
            try {
                String content = response.body().string();
                LOGGER.debug("AIChatEvent[{}]请求结束:{}", event.getId(), content);
                if (response.code() != 200) {
                    //有异常情况
                    handleFail(event, resource, response.code(), content, null);
                } else {
                    AIMessage message = handleChatSuccess(event, content, "message", false);
                    event.setResult(message);
                    event.setState(EventState.SUCCESS);

                }
            } catch (Exception e) {
                LOGGER.error("AIChatEvent[{}]处理响应结果发生异常:{}", event.getId(), e);
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                event.complete();
            }
        }, e -> {
            LOGGER.error("AIChatEvent[{}]请求失败:{}", event.getId(), e);
            event.fail(EventErrorType.NET_ERROR, e);
        });
    }

    /**
     * 以stream方式处理事件
     * @param event
     * @param resource
     * @param request
     */
    private static void processEventByStream(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, true);
        request.asyncSendWithSSE(body -> {

            try {
                LOGGER.debug("AIChatEvent[{}]请求成功:{}", event.getId(), body);
                if (STOP_FLAG.equals(body)) {
                    event.setState(EventState.SUCCESS);
                    //结束
                    return;
                }
                handleChatSuccess(event, body, "delta", true);
            } catch (Exception e) {
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                if (STOP_FLAG.equals(body)) {
                    event.complete();
                }
            }


        }, (res, e) -> {
            try {
                String response = null;
                if (res != null && res.body() != null) {
                    response = res.body().string();
                }
                handleFail(event, resource, res != null ? res.code() : 0, response, e);
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        });
    }

    /**
     * stream请求时，返回的内容放在"delta" 字段中，非stream请求，返回的内容放在message 字段中
     * @param event
     * @param responseContent
     * @param contentField
     */
    private static AIMessage handleChatSuccess(AIChatEvent event, String responseContent, String contentField, boolean stream) {
        JSONObject data = JSON.parseObject(responseContent);
        AIMessage message = new AIMessage();
        message.setId(data.getString("id"));
        event.setResultId(message.getId());
        //获取回复
        JSONObject choice = data.getJSONArray("choices").getJSONObject(0);
        if (choice.containsKey(contentField)) {
            message.setContent(choice.getJSONObject(contentField).getString("content"));
            message.setThink(choice.getJSONObject(contentField).getString("reasoning_content"));
            event.appendResponse(message.getContent());
        } else {
            message.setContent("");
            event.appendResponse("");
        }
        message.setRole(AIMessageRole.assistant);
        event.setResultRole(message.getRole());
        String finishReason = choice.getString("finish_reason");
        if (finishReason != null) {
            switch (finishReason) {
                case "stop":
                    event.setFinishReason(MessageFinishReason.STOP);
                    break;
                case "length":
                    event.setFinishReason(MessageFinishReason.LENGTH_LIMIT);
                    break;
                case "content_filter":
                    event.setFinishReason(MessageFinishReason.CONTENT_FILTER);
                    break;
                case "tool_calls":
                    event.setFinishReason(MessageFinishReason.TOOL_CALL);
                    break;
                default:
                    event.setFinishReason(MessageFinishReason.OTHER);
            }
        }
        //获取使用情况
        data = data.getJSONObject("usage");
        if (data != null) {
            event.setPromptTokenCount(data.getInteger("prompt_tokens"));
            event.setTotalTokenCount(data.getInteger("total_tokens"));
            event.setCompletionTokenCount(data.getInteger("completion_tokens"));
        }

        if (StrUtil.isNotBlank(message.getThink()) ) {
            processThinking(event, stream, message.getThink());
        } else if (message.getThink() == null && event.getThinkStatus() == ThinkStatus.RUNNING) {
            event.thinkOver("");
        }
        if (message.getContent() != null) {
            event.success(message);
        }
        return message;
    }

    /**
     * 处理深度思考内容
     * @param event
     * @param stream
     * @param think
     */
    private static void processThinking(AIChatEvent event, boolean stream, String think) {
        if (stream == false) {
            //非流式
            event.thinkOver(think);
            return;
        }
        //流式
        event.thinking(think);
    }

    /**
     * 失败时的处理
     * @param event
     * @param resource
     * @param code
     * @param responseContent
     * @param exception
     */
    private static void handleFail(AIEvent event, AIResource resource, int code, String responseContent, Throwable exception) {
        JSONObject error  = null;;
        LOGGER.error("AIEvent[{}]请求失败:{}_{},异常：{}", event.getId(), code, responseContent, exception);
        if (exception == null) {
            //ai服务异常
            switch (code) {
                case 400:{
                    error = JSON.parseObject(responseContent);
                    event.fail(EventErrorType.REQUEST_PARAM_ERROR, responseContent);
                    break;
                }

                case 401:
                    event.fail(EventErrorType.API_KEY_INVALID, responseContent);
                    break;
                case 429:{
                    event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, responseContent);
                    break;
                }
                case 500:
                    event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                    break;
                default:
                    //未知错误
                    event.fail(EventErrorType.UNKNOWN,responseContent);
                    break;
            }
        } else {
            //程序异常
            event.fail(EventErrorType.NET_ERROR, exception);
        }
    }




    private RequestInstance buildRequest(AIChatEvent event,  AIResource resource) {
        //资源指定了apiurl就用资源的，否则用默认的
        String apiUrl = StrUtil.isNotBlank(resource.getApiUrl()) ? resource.getApiUrl() : API_URL;
        RequestInstance request = HttpUtil.post(resource.getClient(), apiUrl, event);
        request.setAuthorization(resource.getApiKey());
        if (resource.getAppId() != null) {
            request.addHeader("appid", resource.getAppId());
        }
        request.addJsonParam(AIParamKey.MODEL, event.getModelName());
        if (event.getTemperature() != null) {
            request.addJsonParam(AIParamKey.TEMPERATURE, event.getTemperature());
        }
        if (event.getTopP() != null) {
            request.addJsonParam(AIParamKey.TOP_P, event.getTopP());
        }
        if (event.getMaxTokens() != null) {
            request.addJsonParam("max_completion_tokens", event.getMaxTokens());
        }
        if (event.isThinkingEnable() != null) {
            if (event.isThinkingEnable()) {
                request.addJsonParam(AIParamKey.ENABLE_THINKING, true);
                if (event.getThinkingBudget() != null) {
                    request.addJsonParam(AIParamKey.THINKING_BUDGET, event.getThinkingBudget());
                }
            } else {
                request.addJsonParam(AIParamKey.ENABLE_THINKING, false);
            }
        }

        if (event.getParams() != null) {
            request.addAllJsonParam(event.getParams());
        }

        if (event.getMessages() != null) {
            JSONArray messages = new JSONArray(event.getMessages().size());
            JSONObject item = null;
            List<AIMessage> list = event.getMessages();
            for (AIMessage message : list) {
                item = new JSONObject();
                switch (message.getRole()) {
                    case user:
                        item.put("role", "user");
                        break;
                    case assistant:
                        item.put("role", "assistant");
                        break;
                    case system:
                        item.put("role", "system");
                        break;
                    default:
                        continue;
                }
                if (message.getImgUrl() == null) {
                    item.put("content", message.getContent());
                } else {
                    //带图片
                    JSONArray content = new JSONArray();
                    content.add(new JSONObject().fluentPut("type", "text").fluentPut("text", message.getContent()));
                    for (String img : message.getImgUrl()) {
                        content.add(new JSONObject().fluentPut("type", "image_url").fluentPut("image_url", img));
                    }
                    item.put("content", content);
                }
                messages.add(item);
            }
            request.addJsonParam(AIParamKey.MESSAGES, messages);
        }
        return request;
    }


    @Override
    public void handleCancel(AIChatEvent event, AIResource resource) {
        LOGGER.warn("{}事件取消执行", event.getId());
        RequestInstance requestInstance = event.getAttr(EventAttrKey.REQUEST_INSTANCE);
        requestInstance.cancal();
    }
}
